package cn.xw.utils.security.rsa;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import sun.security.x509.*;

import java.io.*;
import java.math.BigInteger;
import java.net.URL;
import java.nio.file.Files;
import java.security.*;
import java.security.cert.CertificateEncodingException;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.*;

/**
 * 密钥对生成工具，会生成如：密钥库、私钥、公钥、数字证书
 * 使用前请先构建对象，
 * 如：自定义构建 KeyPairUtils(RSAProperties rsaPro)
 * 如：使用默认方式：KeyPairUtils()
 *
 * @author Anhui AntLaddie（博客园蚂蚁小哥）
 * @version 1.0
 */
public class KeyPairUtils {

    private static final Logger log = LoggerFactory.getLogger(KeyPairUtils.class);

    // 配置一些公共信息，若不需要默认配置则可以自己修改
    private RSAProperties rsaPro = new RSAProperties();

    public KeyPairUtils() {
    }

    public KeyPairUtils(RSAProperties rsaPro) {
        this.rsaPro = rsaPro;
    }

    /***
     * 生成密钥对信息，生成的文件存储到指定的类路径下
     */
    public void generateKeyPair() {
        // 初始化密钥对信息
        KeyPair keyPair = null;
        try {
            // 构建需要生成文件的路径文件夹
            buildWritePath();
            // 使用指定算法，生成一个指定位的密钥对。并设置签名的取模位数，生成密钥对
            KeyPairGenerator keyGen = KeyPairGenerator.getInstance(RSAProperties.KEY_ALGORITHM);
            keyGen.initialize(rsaPro.getModulus());
            keyPair = keyGen.generateKeyPair();

            // 根据传入公钥私钥（密钥对）生成自签名的数字证书
            X509Certificate selfSignedCertificate = createSelfSignedCertificate(keyPair);
            // 输出数字证书、公钥和私钥的 Base64 编码值
            String certificate = Base64.getEncoder().encodeToString(selfSignedCertificate.getEncoded());
            fileWrite(RSAProperties.DIGITAL_CERT, certificate);
            // 写出公钥信息
            String publicKey = Base64.getEncoder().encodeToString(keyPair.getPublic().getEncoded());
            fileWrite(RSAProperties.PUB_FILE_NAME, publicKey);

            // 判断为简单模式还是标准模式（标准模式不可以生成私钥文件，它存在密钥库中）
            if (!rsaPro.getSimpleGeneration()) {
                // 写出私钥信息
                String privateKey = Base64.getEncoder().encodeToString(keyPair.getPrivate().getEncoded());
                fileWrite(RSAProperties.PRI_FILE_NAME, privateKey);
            }
            // 若不需要可以注释，或者更改日志级别
            //log.info("==========>日志记录<=========");
            //log.info("==>数字证书：\r\n        {}", certificate);
            //log.info("==>公钥信息：\r\n        {}", publicKey);
            //log.info("==>私钥信息：\r\n        {}", privateKey);
        } catch (NoSuchAlgorithmException e) {
            log.info("加密算法或消息摘要算法不存在：{}", e.getMessage());
            throw new RuntimeException(e);
        } catch (CertificateEncodingException e) {
            log.info("证书编码或解码过程发生异常：{}", e.getMessage());
            throw new RuntimeException(e);
        }
    }

    /***
     * 根据传入的密钥对 生成自签名证书，并返回该证书对象
     * @param keyPair 密钥对，包括公钥和私钥
     * @return 生成的自签名证书对象
     */
    private X509Certificate createSelfSignedCertificate(KeyPair keyPair) {
        // 获取当前时间和证书有效期
        Calendar calendar = Calendar.getInstance();
        Date currentTime = calendar.getTime();
        calendar.add(Calendar.MINUTE, rsaPro.getCertificateExpTime());
        // 初始化证书信息
        X509Certificate x509Certificate = null;
        try {
            // ========================== 构建证书信息 (Start) ==========================
            X509CertInfo info = new X509CertInfo();
            X500Name subject = new X500Name("C=" + rsaPro.getCountry() + ", ST=" + rsaPro.getProvince()
                    + ", L=" + rsaPro.getCity() + ", " + "O=" + rsaPro.getOrganization() + ", OU="
                    + rsaPro.getDepartment() + ", CN=" + rsaPro.getCompany());
            info.set(X509CertInfo.VERSION, new CertificateVersion(CertificateVersion.V3));  // 证书版本
            info.set(X509CertInfo.SUBJECT, subject);                                        // 证书主题
            info.set(X509CertInfo.ISSUER, subject);                                         // 颁发者
            info.set(X509CertInfo.KEY, new CertificateX509Key(keyPair.getPublic()));        // 公钥
            info.set(X509CertInfo.SERIAL_NUMBER, new CertificateSerialNumber(               // 序列号
                    new BigInteger(64, new SecureRandom())));
            info.set(X509CertInfo.VALIDITY, new CertificateValidity(currentTime, calendar.getTime())); //有效期
            AlgorithmId algorithm = new AlgorithmId(AlgorithmId.sha256WithRSAEncryption_oid);   // 算法标识
            info.set(X509CertInfo.ALGORITHM_ID, new CertificateAlgorithmId(algorithm));
            // 签名证书信息
            X509CertImpl certImpl = new X509CertImpl(info);
            // ========================== 构建证书信息 (End) ==========================
            // 获取私钥信息
            PrivateKey aPrivate = keyPair.getPrivate();
            // 数字证书签名，并设置到X509证书类里
            //      SHA256WithRSA参数说明：使用的是SHA-256算法生成摘要，并使用RSA算法对该摘要进行签名。
            certImpl.sign(aPrivate, "SHA256WithRSA");
            x509Certificate = certImpl;
            // 判断为简单模式还是标准模式（标准模式存在密钥库）
            if (rsaPro.getSimpleGeneration()) {
                //往文件中写出个密钥库（密钥库（证书+私钥））
                writePrivateKey(x509Certificate, keyPair);
            }
            // 异常处理
        } catch (IOException e) {
            log.info("数字证书保存或加载异常：{}", e.getMessage());
        } catch (CertificateException e) {
            log.info("数字证书在处理或解析过程中发生了一般性错误：{}", e.getMessage());
        } catch (NoSuchAlgorithmException e) {
            log.info("数字证书签名时，使用了系统不支持的算法：{}", e.getMessage());
        } catch (InvalidKeyException e) {
            log.info("无效密钥：{}", e.getMessage());
        } catch (NoSuchProviderException e) {
            log.info("加密服务提供程序不存在：{}", e.getMessage());
        } catch (SignatureException e) {
            log.info("数字签名验证失败：{}", e.getMessage());
        } catch (KeyStoreException e) {
            log.info("密钥库操作失败：{}", e.getMessage());
        }
        // 若生成失败为null，异常结束程序
        if (x509Certificate == null) {
            throw new RuntimeException("数字证书签名失败，请检查后重试...");
        }
        // 其实上面对生成的信息写出到密钥库（KeyStore）里，它是一个存储密钥和证书的安全存储区域，通常以文件的形式存在。
        // 若想操作那个密钥库信息则需要KeyStore操作；而这里方法返回的是数字证书对象（X509Certificate），而非证书信息。
        // 数字证书对象内包含证书信息、公钥信息、密钥信息；
        return x509Certificate;
    }

    /***
     * 生成密钥库（数字证书+私钥）
     * @param x509Certificate  数字证书信息
     * @param keyPair 公钥私钥信息
     */
    private void writePrivateKey(X509Certificate x509Certificate, KeyPair keyPair) throws KeyStoreException,
            CertificateException, IOException, NoSuchAlgorithmException {
        // 获取要写入密钥库的路径
        String path = Objects.requireNonNull(this.getClass().getClassLoader()
                .getResource(RSAProperties.WRITE_PATH + rsaPro.getDatePath())).getPath();
        // 创建一个空的密钥库对象，并指定对应密钥库类型
        KeyStore keyStore = KeyStore.getInstance(rsaPro.getKeyStoreType());
        keyStore.load(null);
        // 将数字证书存储到KeyStore中，并为其设置别名，(若存在多个证书则可以多次setCertificateEntry)
        keyStore.setCertificateEntry(rsaPro.getCertificateAlias(), x509Certificate);
        // 密码处理（密钥库密码已经私钥密码处理）
        char[] keyStorePwd = rsaPro.getKeyStorePassword().toCharArray();
        char[] priPwd = rsaPro.getPrivateKeyPassword().toCharArray();
        // ==================== 写出私钥和公用密钥库 ====================
        // 将一个私钥和相应的证书链存储到密钥库中（就是私钥存储到密钥库里，并设置上密码）
        //     参数1：指定私钥的别名。该别名必须是唯一的，用于在密钥库中标识特定的私钥和证书链。
        //     参数2：指定要存储的私钥。通常情况下，这个私钥是一个 RSA 或 DSA 密钥对
        //     参数3：指定私钥的保护密码。这个密码用于保护存储在密钥库中的私钥
        //     参数4：指定证书链。该证书链应该包括服务器证书以及它所依赖的任何中间证书和根证书
        // 注意：这里的key既可以放公钥也可以放私钥
        keyStore.setKeyEntry(rsaPro.getPriAlias(), keyPair.getPrivate(), priPwd,
                new X509Certificate[]{x509Certificate});
        // 写出密钥库信息文件（密钥库并指定密码保护）
        try (FileOutputStream out = new FileOutputStream(path + rsaPro.getKeystoreInfoFileName())) {
            keyStore.store(out, keyStorePwd);
        }
    }

    /***
     * 把数据写入到指定的文件中
     * @param fileName 文件名称（固定的三个文件名称）
     * @param data 要写入的数据
     */
    private void fileWrite(String fileName, String data) {
        // 获取要写入的路径信息
        String path = Objects.requireNonNull(this.getClass().getClassLoader()
                .getResource(RSAProperties.WRITE_PATH + rsaPro.getDatePath())).getPath();
        File file = new File(path + fileName);
        // 根据文件名称判断当前写入的是什么文件信息（因为文件名称是不能修改的）
        String tagInfo =
                fileName.equals(RSAProperties.PUB_FILE_NAME) ? " PUBLIC KEY" :
                        fileName.equals(RSAProperties.PRI_FILE_NAME) ? " PRIVATE KEY" :
                                fileName.equals(RSAProperties.DIGITAL_CERT) ? " DIGITAL CERTIFICATE" : "";
        // 对数据切割分片
        final int LINE_LENGTH = 64;
        StringBuilder formattedKey = new StringBuilder();
        int offset = 0;
        while (offset < data.length()) {
            int endIndex = Math.min(offset + LINE_LENGTH, data.length());
            formattedKey.append(data, offset, endIndex);
            formattedKey.append("\n");
            offset = endIndex;
        }
        try {
            // 构建输出流并写出文件信息
            BufferedWriter bufferedWriter =
                    new BufferedWriter(new OutputStreamWriter(Files.newOutputStream(file.toPath())));
            String start = "-----BEGIN " + tagInfo + "-----";
            String end = "-----END " + tagInfo + "-----";
            bufferedWriter.write(start + "\n", 0, start.length() + 1);
            bufferedWriter.write(formattedKey.toString(), 0, formattedKey.length());
            bufferedWriter.write(end, 0, end.length());
            bufferedWriter.flush();
        } catch (IOException e) {
            log.info("生成的公钥、私钥、数字证书失败：{}", e.getMessage());
        }
    }

    /***
     * 构建路径信息是否存在（classpath:security/rsa/日期yyyyMMddHH）具体路径看writePath配置；
     * 若存在则不操作，反之创建
     */
    private void buildWritePath() {
        // 获取具体路径，若不存在则返回null
        URL urlPath = this.getClass().getClassLoader()
                .getResource(RSAProperties.WRITE_PATH + rsaPro.getDatePath());
        // 代表路径不存在则创建信息（反之程序不处理）
        if (urlPath == null) {
            // 获取根路径信息
            String path = Objects.requireNonNull(this.getClass().getClassLoader().getResource(""))
                    .getPath();
            // 创建文件信息
            File file = new File(path + RSAProperties.WRITE_PATH + rsaPro.getDatePath());
            if (!file.exists()) {
                // 创建文件夹信息
                boolean mkdirs = file.mkdirs();
                if (!mkdirs) {
                    throw new RuntimeException("创建存放密钥库、密钥、公钥文件夹失败....");
                }
            }
        }
    }
}
